import requests
import threading, time, os
import aiohttp
from aiohttp import ClientWebSocketResponse
import numpy as np
from enum import Enum
from typing import AsyncGenerator, Generator
import asyncio
from contextlib import asynccontextmanager, contextmanager

# 文本(txt)的内容输入源
class TxtInputSource:
    """
    TTS时的内容输入源
    """
    def __init__(self, filepath:str):
        """
        初始化. 注意: 若文件中行过多(>100)且处理任务用时很小/无需休眠等待(<1s)不要开启异步模式
        Args:
            filepath (str): 内容源的路径
        """
        self.iter = self.__read_file(filepath)
    
    @staticmethod
    def tail_symbol()->str:
        """
        内容源结束的标志
        Returns:
            str: 结束的标志
        """
        return 'EOF'

    # 读取文件源行的生成器
    def __read_file(self, filepath:str):
        """
        读取文件源行的同步生成器
        Args:
            filepath (str): 内容源的路径

        Yields:
            _type_: 行内容生成器
        """
        self.__file_exists_check(filepath)
        with open(filepath, encoding='utf-8') as fp:
            for line in fp:
                yield line.strip()  
                
    def __file_exists_check(self, filepath:str)->None:
        assert os.path.exists(filepath), f'file:{filepath} not find'
        
    # 取出下一行
    # 若文件结束, 则返回None
    def poll_segment(self)->str:
        """
        取出下一行的内容. 适用于: 同步模式
        Returns:
            str: 文件内容. 无可用时返回None
        """
        assert isinstance(self.iter, Generator), 'current is not Sync Generator type'
        return next(self.iter, None)
    
    def close(self):
        if isinstance(self.iter, Generator):
            self.iter.close() 
        if isinstance(self.iter, AsyncGenerator):
            asyncio.run_coroutine_threadsafe(self.iter.aclose(), loop=asyncio.get_running_loop())
            
    @property
    def sync_gen(self):
        assert isinstance(self.iter, Generator), 'current is not Sync Generator type'
        return self.iter


class RecognitionTransType(Enum):
    WAV = 'wav'
    NP = 'numpy'
    NAN = ''


class TraineRestConnection:
    """
    RestAPI会话连接
    """

    def __init__(self, auth_endpoint: str, rest_email: str, rest_password: str, timeout: int = 60):
        """
        初始化
        Args:
            auth_endpoint (str): Rest的会话连接地址
            rest_email (str): Rest的会话连接需要的帐号
            rest_password (str): Rest的会话连接需要的密码
            timeout (int, optional): 连接的超时时长(秒). Defaults to 60.
        """
        self.__REST_AUTH_ENDPOINT = auth_endpoint
        self.__REST_EMAIL = rest_email
        self.__REST_PSWD = rest_password
        self.__connect_timeout = timeout
        #
        # access_token有效期为: 30分钟/半小时
        self.__access_token: str = None
        self.__query_authtoken_thread = None
        self.__authtoken_thread_quit = None
        #
        self._header = {'User-Agent': 'Mozilla/5.0', "Content-Type": "application/x-www-form-urlencoded"}

    @property
    def auth_api_uri(self) -> str:
        return self.__REST_AUTH_ENDPOINT

    @property
    def access_token(self) -> str:
        _token = self.__access_token
        return _token

    def start(self) -> None:
        """
        开始连接
        """

        def query_auth_token(cls: TraineRestConnection, quitEvent: threading.Event):
            while not quitEvent.is_set():
                try:
                    cls.auth_passport()
                    time.sleep(28 * 60.0)
                except InterruptedError:
                    break

        if self.__query_authtoken_thread is None:
            self.__authtoken_thread_quit = threading.Event()
            self.__query_authtoken_thread = threading.Thread(
                name="query-train-authtoken-listener",
                target=query_auth_token,
                args=(
                    self,
                    self.__authtoken_thread_quit,
                ),
                daemon=True
            )
            self.__query_authtoken_thread.start()

    def __wait_connected(self, timeout: int = 60) -> None:
        max_wait_sec = timeout
        while max_wait_sec > 0:
            if self.__access_token is None:
                max_wait_sec -= 1
                time.sleep(1.0)
                print('lost 1 second')
            else:
                break

        if max_wait_sec == 0:
            raise Exception

    # OAuth2PasswordRequestForm
    def auth_passport(self) -> None:
        """
        执行REST-API身份验证

        Raises:
            InterruptedError: 验证失败抛出
        """
        reqURL = '{}/{}'.format(self.__REST_AUTH_ENDPOINT, 'login')
        payload = {'username': self.__REST_EMAIL, 'password': self.__REST_PSWD, }
        response: requests.Response = requests.post(reqURL, headers=self._header, data=payload)
        resJson = (response.json())  # {'access_token': ?, 'refresh_token': ?}
        if response.ok:
            self.__access_token = resJson['access_token']

        raise InterruptedError

    # SystemUser
    def show_passport(self) -> dict[str, str]:
        """
        显示REST-API身份信息

        Returns:
            dict[str,str]: {'id': ?, 'email': ?, }
        """
        reqURL = '{}/{}'.format(self.__REST_AUTH_ENDPOINT, 'me')
        _headers = {'User-Agent': 'Mozilla/5.0', 'Authorization': f'Bearer {self.access_token}'}
        response: requests.Response = requests.get(reqURL, headers=_headers)
        resJson = (response.json())  # {'id': ?, 'email': ?, }
        if response.ok:
            return dict(resJson)
        else:
            return None

    def close(self) -> None:
        """
        关闭连接
        """
        if self.__query_authtoken_thread:
            self.__authtoken_thread_quit.set()
            print('wait auth thread close')
            self.__query_authtoken_thread.join(timeout=60.0)
            self.__query_authtoken_thread = None

    def __enter__(self):
        self.start()
        self.__wait_connected(self.__connect_timeout)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class TraineRestClient:
    """
    推理服务RestAPI客户端
    """

    def __init__(self, api_endpoint: str, auth_endpoint: str, rest_email: str, rest_password: str):
        """
        初始化
        Args:
            api_endpoint (str): Rest的API地址
            auth_endpoint (str): Rest的会话连接地址
            rest_email (str): Rest的会话连接需要的帐号
            rest_password (str): Rest的会话连接需要的密码
        """
        self.__REST_API_ENDPOINT = api_endpoint
        self.__REST_AUTH_ENDPOINT = auth_endpoint
        self.__REST_EMAIL = rest_email
        self.__REST_PSWD = rest_password
        self._header = {'User-Agent': 'Mozilla/5.0', "Content-Type": "application/x-www-form-urlencoded"}
        #
        self.__connect = None

    @property
    def rest_api_uri(self) -> str:
        return self.__REST_API_ENDPOINT

    @contextmanager
    def connection(self, timeout: int = 60):
        """
        RestAPI会话连接生成器. 适用于方法的参数中存在: connect:TraineRestConnection
        Args:
            timeout (int, optional): _description_. Defaults to 60.

        Yields:
            _type_: _description_
        """
        with TraineRestConnection(self.__REST_AUTH_ENDPOINT, self.__REST_EMAIL, self.__REST_PSWD, timeout) as connect:
            yield connect

    def start(self):
        """
        初始化RestAPI会话连接
        """
        if self.__connect is None:
            self.__connect = TraineRestConnection(self.__REST_AUTH_ENDPOINT, self.__REST_EMAIL, self.__REST_PSWD)
            self.__connect.start()

    def translate(self, prompt: str, language: str) -> str | None:
        """
        执行REST-API 翻译
        注意: 需要提前调用start方法完成TraineRestConnection连接
        Args:
            prompt (str): 待翻译的文本
            language (str): 翻译后的目标语种, 支持: English|Thai|Arabic

        Returns:
            str|None: _description_
        """
        assert self.__connect, 'TraineRestConnection instance is None'
        #
        reqURL = '{}/{}/{}'.format(self.__REST_API_ENDPOINT, 'translation', 'nllb')
        _headers = {**self._header, **{'Authorization': f'Bearer {self.__connect.access_token}'}}
        payload = {'prompt': prompt, 'language': language}
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload)
        resJson = (response.json())  # {"result": ?, "language": language}
        return resJson['result'] if response.ok else None

    def speech(self, prompt: str, language: str) -> tuple[str, str] | None:
        """
        执行REST-API 文本语音生成
        注意: 需要提前调用start方法完成TraineRestConnection连接
        Args:
            prompt (str): 待生成语音的文本
            language (str): 文本所属的语言, 支持: English|Thai|Arabic

        Returns:
            tuple[str,str]|None: tuple[0] 语音可访问的URL, tuple[1] 文件名
        """
        assert self.__connect, 'TraineRestConnection instance is None'
        #
        reqURL = '{}/{}'.format(self.__REST_API_ENDPOINT, 'speech')
        _headers = {**self._header, **{'Authorization': f'Bearer {self.__connect.access_token}'}}
        payload = {'prompt': prompt, 'language': language}
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload)
        resJson = (response.json())  # {"text": prompt, "language": language, "path": ?, "filename": ?}
        if response.ok:
            return resJson['path'], resJson['filename']
        else:
            return None

    def batch_translate(self, prompts: list[str], language: str) -> list[dict]:
        """
        执行REST-API 批量翻译
        注意: 需要提前调用start方法完成TraineRestConnection连接
        Args:
            prompts (list[str]): 待翻译的文本列表.由逗号分隔
            language (str): 翻译后的目标语种, 支持: English|Thai|Arabic

        Returns:
            list[dict]: 翻译后的结果列表. 结构: [{prompt:输入内容, result:翻译后的内容}, ...]
        """
        assert self.__connect, 'TraineRestConnection instance is None'
        #
        reqURL = '{}/{}/{}'.format(self.__REST_API_ENDPOINT, 'translation/batch', 'nllb')
        _headers = {**self._header, **{'Authorization': f'Bearer {self.__connect.access_token}'}}
        payload = {'content': ','.join(prompts), 'language': language}
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload)
        resJson = (response.json())
        # ---------------------------------------------------------------
        # { result: [
        #     {prompt:x, result:'?'},
        #     {prompt:x, result:'?'},
        #     ...
        #   ],
        #   language: '?'
        # }
        return resJson['result'] if response.ok else None

    def recognition(self, audio_file_path: str) -> str | None:
        """
        执行REST-API 语音识别
        注意: 需要提前调用start方法完成TraineRestConnection连接
        Args:
            audio_file_path (str): 音频文件地址

        Returns:
            str|None: _description_
        """
        assert self.__connect, 'TraineRestConnection instance is None'
        #
        reqURL = '{}/{}'.format(self.__REST_API_ENDPOINT, 'recognition')
        _headers = {**self._header,
                    **{"Content-Type": "multipart/form-data", 'Authorization': f'Bearer {self.__connect.access_token}'}}
        payload = {'model': 'faster-whisper'}
        upload_files = {'audio': open(audio_file_path, 'rb')}
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload, files=upload_files)
        resJson = (response.json())  # {"result": ?}
        return resJson['result'] if response.ok else None

    def create_chat(self, language: str, ttl: int = 3600) -> str | None:
        """
        执行REST-API 创建流式(文本翻译+语音生成)交互Token
        注意: 需要提前调用start方法完成TraineRestConnection连接
        Args:
            language (str): 生成语音的目标语种, 支持: English|Thai|Arabic
            ttl (int, optional): 交互有效时长. Defaults to 3600.

        Returns:
            str|None: 非None为(/stream/chat的)通信Token
        """
        assert self.__connect, 'TraineRestConnection instance is None'
        #
        reqURL = '{}/{}'.format(self.__REST_API_ENDPOINT, 'chat/create')
        _headers = {**self._header, **{'Authorization': f'Bearer {self.__connect.access_token}'}}
        payload = {'language': language, 'ttl': ttl, 'translation': 'nllb', 'speech': 'mms'}
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload)
        resJson = (response.json())  # {'result': ?, 'expire': ttl}
        if response.ok:
            return resJson['result']
        else:
            return None

    def create_speech(self, language: str, connect: TraineRestConnection, sample_rate: int = 16000, channel_count: int = 1, ttl: int = 3600) -> str | None:
        """
        执行REST-API 创建流式(文本到语音)交互Token
        Args:
            language (str): 生成语音的目标语种
            connect (TraineRestConnection): RestAPI会话连接
            sample_rate (int, optional): 音频采样率. Defaults to 16000.
            channel_count (int, optional): 音频声道数量. Defaults to 1.
            ttl (int, optional): 交互有效时长. Defaults to 3600.

        Returns:
            str|None: 非None为(/stream/speech的)通信Token
        """
        reqURL = '{}/{}'.format(self.__REST_API_ENDPOINT, 'speech/create')
        _headers = {**self._header, **{'Authorization': f'Bearer {connect.access_token}'}}
        payload = {'language': language, 'ttl': ttl, 'model': 'mms', 'rate': sample_rate, 'channel': channel_count}
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload)
        resJson = (response.json())  # {'result': ?, 'expire': ttl}
        if response.ok:
            return resJson['result']
        else:
            return None

    def create_recognition(self, transfer_type: RecognitionTransType, connect: TraineRestConnection, ttl: int = 3600, ) -> str | None:
        """
        执行REST-API 创建流式(语音识别)交互Token
        Args:
            transfer_type(RecognitionTransType): 传输的二进制类型
            connect (TraineRestConnection): RestAPI会话连接
            ttl (int, optional): 交互有效时长. Defaults to 3600.

        Returns:
            str|None: _description_
        """
        reqURL = '{}/{}'.format(self.__REST_API_ENDPOINT, 'recognition/create')
        _headers = {**self._header, **{'Authorization': f'Bearer {connect.access_token}'}}
        payload = {'ttl': ttl, 'model': 'faster-whisper', 'type': transfer_type.value, }
        response: requests.Response = requests.post(reqURL, headers=_headers, data=payload)
        resJson = (response.json())  # {'result': ?, 'expire': ttl}
        if response.ok:
            return resJson['result']
        else:
            return None

    @asynccontextmanager
    async def chat_stream_pipe(self, token: str) -> AsyncGenerator[ClientWebSocketResponse, None]:
        """
        执行REST-API 文本翻译+生成语音(wav_url)流.
        注意: 需要提前调用start方法完成TraineRestConnection连接
        Args:
            token (str): 通信标识

        Returns:
            AsyncGenerator[ClientWebSocketResponse, None]: _description_

        Yields:
            Iterator[AsyncGenerator[ClientWebSocketResponse, None]]: _description_

        Usage::
            >>> async with client.chat_stream_pipe(token):
            >>>     pass:

        """
        assert self.__connect, 'TraineRestConnection instance is None'
        #
        async with aiohttp.ClientSession() as session:
            # 连接流式交互的WebSocket地址
            # REST-API Auth
            reqURL = '{}/{}?token={}'.format(self.__REST_API_ENDPOINT, 'stream/chat', token)
            _headers = {'Authorization': f'Bearer {self.__connect.access_token}'}
            async with session.ws_connect(reqURL, headers=_headers) as rest_ws:
                yield rest_ws
                await rest_ws.close()

    async def speech_stream_pipe(self, input_file_path: str, connect: TraineRestConnection, token: str) -> AsyncGenerator[tuple[np.ndarray, str], None]:
        """
        执行REST-API 文本生成语音(numpy)流
        Args:
            input_file_path (str): 内容源文件地址
            connect (TraineRestConnection): RestAPI会话连接
            token (str): 通信标识

        Raises:
            Exception: 中断异常/标志通信结束

        Returns:
            AsyncGenerator[tuple[np.ndarray, str], None]: dtype=int16的numpy ndarray数组

        Yields:
            Iterator[AsyncGenerator[tuple[np.ndarray, str], None]]: _description_

        Usage::
            >>> import sounddevice
            >>> async for nparr, word in client.speech_stream_pipe(filepath, connect, token):
            >>>     try:
            >>>         print('word: %s, numpy len:%d' % (word, len(nparr)))
            >>>         sounddevice.play(data=nparr, samplerate=16000, blocking=True)
            >>>         await asyncio.sleep(0.1)
            >>>     except Exception:
            >>>         break
        """
        reqURL = '{}/{}?token={}'.format(self.__REST_API_ENDPOINT, 'stream/speech', token)
        _headers = {'User-Agent': 'Mozilla/5.0', 'Authorization': f'Bearer {connect.access_token}'}
        #
        tis = TxtInputSource(input_file_path)
        async with aiohttp.ClientSession() as session:
            async with session.ws_connect(reqURL, headers=_headers) as rest_ws:
                while True:
                    word = tis.poll_segment()
                    if word is None:
                        break
                    await rest_ws.send_str(word)
                    result = await rest_ws.receive_bytes()
                    if len(result) > 0:
                        nparr: np.ndarray = np.frombuffer(result, dtype=np.int16)
                        yield nparr, word
                    else:
                        raise Exception

    @asynccontextmanager
    async def recognition_stream_pipe(self, connect: TraineRestConnection, token: str) -> AsyncGenerator[ ClientWebSocketResponse, None]:
        """
        执行REST-API 语音识别流
        Args:
            connect (TraineRestConnection): RestAPI会话连接
            token (str): 通信标识

        Returns:
            AsyncGenerator[ClientWebSocketResponse, None]: _description_

        Yields:
            Iterator[AsyncGenerator[ClientWebSocketResponse, None]]: _description_

        Usage::
            >>> async with client.recognition_stream_pipe(connect, token):
            >>>     pass:
        """
        async with aiohttp.ClientSession() as session:
            # 连接流式交互的WebSocket地址
            # REST-API Auth
            reqURL = '{}/{}?token={}'.format(self.__REST_API_ENDPOINT, 'stream/recognition', token)
            _headers = {'Authorization': f'Bearer {connect.access_token}'}
            async with session.ws_connect(reqURL, headers=_headers) as rest_ws:
                yield rest_ws
                await rest_ws.close()

    # UserAuth
    def create_passport(self, usermail: str, userpswd: str) -> str | None:
        """
        创建REST-API 身份信息
        Args:
            usermail (str): 邮箱
            userpswd (str): 密码

        Returns:
            str|None: 失败返回None, 成功返回身份ID
        """
        reqURL = '{}/{}'.format(self.__REST_AUTH_ENDPOINT, 'signup')
        _headers = {**self._header, **{"Content-Type": "application/json"}}
        payload = {'password': userpswd, 'email': usermail}
        response: requests.Response = requests.post(reqURL, headers=_headers, json=payload)
        resJson = (response.json())  # {'id': ?, 'email': ?}
        if response.ok:
            return resJson['id']
        else:
            #
            # 400
            # { "detail": "User with this email already exist"}
            #
            return None

    def close(self) -> None:
        if self.__connect is not None:
            self.__connect.close()